from typing import Literal
from langchain_openai import ChatOpenAI
from langgraph.types import Command
from langgraph.graph import StateGraph, MessagesState, START, END

model = ChatOpenAI() # 온프레미스 환경에서는 자체 호스팅 LLM으로 대체 가능

def supervisor(state: MessagesState) -> Command[Literal["agent_1", "agent_2", END]]:
    response = model.invoke(state["messages"])
    # 다음에 호출할 에이전트 결정
    return Command(goto=response["next_agent"])

def agent_1(state: MessagesState) -> Command[Literal["supervisor"]]:
    response = model.invoke(state["messages"])
    return Command(
        goto="supervisor",
        update={"messages": [response]}
    )

def agent_2(state: MessagesState) -> Command[Literal["supervisor"]]:
    response = model.invoke(state["messages"])
    return Command(
        goto="supervisor",
        update={"messages": [response]}
    )

# 상태 그래프 구성
builder = StateGraph(MessagesState)
builder.add_node(supervisor)
builder.add_node(agent_1)
builder.add_node(agent_2)
builder.add_edge(START, "supervisor")
supervisor_graph = builder.compile()